#include "algorithms.h"
#include <algorithm>

static std::vector<Pos> backtrack(const Grid& g, const std::map<long long,long long>& parent, Pos end) {
    std::vector<Pos> path;
    long long cur = key(end, g.cols);
    auto it = parent.find(cur);
    if (it==parent.end() && !(end.r==g.start.r && end.c==g.start.c)) return path;
    while (true) {
        path.push_back(Pos{(int)(cur / g.cols), (int)(cur % g.cols)});
        if (cur == key(g.start, g.cols)) break;
        cur = parent.at(cur);
    }
    std::reverse(path.begin(), path.end());
    return path;
}

int heuristic(const Grid& g, Pos a, Pos b) {
    int dr = std::abs(a.r - b.r);
    int dc = std::abs(a.c - b.c);
    if (g.diagonal) return std::max(dr, dc);
    return dr + dc;
}

std::vector<Pos> rebuild_path(const Grid& g, const std::map<long long,long long>& parent, Pos end) { return backtrack(g, parent, end); }

std::vector<Pos> rebuild_bi_path(const Grid& g, const std::map<long long,long long>& pa, const std::map<long long,long long>& pb, Pos m) {
    std::vector<Pos> a = backtrack(g, pa, m);
    std::vector<Pos> b;
    long long cur = key(m, g.cols);
    while (cur != key(g.goal, g.cols)) {
        auto it = pb.find(cur);
        if (it==pb.end()) break;
        cur = it->second;
        b.push_back(Pos{(int)(cur / g.cols), (int)(cur % g.cols)});
    }
    a.insert(a.end(), b.begin(), b.end());
    return a;
}

void BFS::reset(const Grid& g) {
    while(!q.empty()) q.pop();
    parent.clear();
    visited.clear();
    frontier_s.clear();
    path_cache.clear();
    fin=false; found=false;
    q.push(g.start);
    visited.insert(key(g.start, g.cols));
    frontier_s.insert(key(g.start, g.cols));
}

void BFS::step(const Grid& g) {
    if (fin) return;
    if (q.empty()) { fin=true; return; }
    Pos u = q.front(); q.pop();
    frontier_s.erase(key(u, g.cols));
    if (u.r==g.goal.r && u.c==g.goal.c) { fin=true; found=true; path_cache = backtrack(g, parent, u); return; }
    for (auto v : g.neighbors(u)) {
        long long kv = key(v, g.cols);
        if (!visited.count(kv)) {
            visited.insert(kv);
            parent[kv] = key(u, g.cols);
            q.push(v);
            frontier_s.insert(kv);
        }
    }
}

AlgoState BFS::state() const { AlgoState s; s.finished=fin; s.found=found; s.visited=visited; s.frontier=frontier_s; s.path=path_cache; return s; }

void DFS::reset(const Grid& g) {
    while(!st.empty()) st.pop();
    parent.clear();
    visited.clear();
    frontier_s.clear();
    path_cache.clear();
    fin=false; found=false;
    st.push(g.start);
    frontier_s.insert(key(g.start, g.cols));
}

void DFS::step(const Grid& g) {
    if (fin) return;
    if (st.empty()) { fin=true; return; }
    Pos u = st.top(); st.pop();
    frontier_s.erase(key(u, g.cols));
    long long ku = key(u, g.cols);
    if (visited.count(ku)) return;
    visited.insert(ku);
    if (u.r==g.goal.r && u.c==g.goal.c) { fin=true; found=true; path_cache = backtrack(g, parent, u); return; }
    auto ns = g.neighbors(u);
    for (auto it = ns.rbegin(); it!=ns.rend(); ++it) {
        long long kv = key(*it, g.cols);
        if (!visited.count(kv)) { parent[kv]=ku; st.push(*it); frontier_s.insert(kv); }
    }
}

AlgoState DFS::state() const { AlgoState s; s.finished=fin; s.found=found; s.visited=visited; s.frontier=frontier_s; s.path=path_cache; return s; }

void Dijkstra::reset(const Grid& g) {
    while(!pq.empty()) pq.pop();
    parent.clear(); dist.clear(); visited.clear(); path_cache.clear(); fin=false; found=false;
    dist[key(g.start,g.cols)] = 0;
    pq.push(Node{g.start,0});
    frontier_s.clear();
    frontier_s.insert(key(g.start, g.cols));
}

void Dijkstra::step(const Grid& g) {
    if (fin) return;
    if (pq.empty()) { fin=true; return; }
    Node cur = pq.top(); pq.pop();
    frontier_s.erase(key(cur.p, g.cols));
    long long ku = key(cur.p, g.cols);
    if (visited.count(ku)) return;
    visited.insert(ku);
    if (cur.p.r==g.goal.r && cur.p.c==g.goal.c) { fin=true; found=true; path_cache = backtrack(g, parent, cur.p); return; }
    for (auto v : g.neighbors(cur.p)) {
        long long kv = key(v, g.cols);
        int nd = dist[ku] + 1;
        if (!dist.count(kv) || nd < dist[kv]) {
            dist[kv] = nd;
            parent[kv] = ku;
            pq.push(Node{v, nd});
            frontier_s.insert(kv);
        }
    }
}

AlgoState Dijkstra::state() const { AlgoState s; s.finished=fin; s.found=found; s.visited=visited; s.frontier=frontier_s; s.path=path_cache; return s; }

void AStar::reset(const Grid& g) {
    while(!pq.empty()) pq.pop(); parent.clear(); dist.clear(); visited.clear(); path_cache.clear(); fin=false; found=false;
    dist[key(g.start,g.cols)] = 0;
    int h = heuristic(g, g.start, g.goal);
    pq.push(Node{g.start,0,h});
    frontier_s.clear();
    frontier_s.insert(key(g.start, g.cols));
}

void AStar::step(const Grid& g) {
    if (fin) return;
    if (pq.empty()) { fin=true; return; }
    Node cur = pq.top(); pq.pop();
    frontier_s.erase(key(cur.p, g.cols));
    long long ku = key(cur.p, g.cols);
    if (visited.count(ku)) return;
    visited.insert(ku);
    if (cur.p.r==g.goal.r && cur.p.c==g.goal.c) { fin=true; found=true; path_cache = backtrack(g, parent, cur.p); return; }
    for (auto v : g.neighbors(cur.p)) {
        long long kv = key(v, g.cols);
        int nd = dist[ku] + 1;
        if (!dist.count(kv) || nd < dist[kv]) {
            dist[kv] = nd;
            int h = heuristic(g, v, g.goal);
            parent[kv] = ku;
            pq.push(Node{v, nd, nd + h});
            frontier_s.insert(kv);
        }
    }
}

AlgoState AStar::state() const { AlgoState s; s.finished=fin; s.found=found; s.visited=visited; s.frontier=frontier_s; s.path=path_cache; return s; }

void GBFS::reset(const Grid& g) {
    while(!pq.empty()) pq.pop(); parent.clear(); visited.clear(); path_cache.clear(); fin=false; found=false;
    pq.push(Node{g.start, heuristic(g, g.start, g.goal)});
    frontier_s.clear();
    frontier_s.insert(key(g.start, g.cols));
}

void GBFS::step(const Grid& g) {
    if (fin) return;
    if (pq.empty()) { fin=true; return; }
    Node cur = pq.top(); pq.pop();
    frontier_s.erase(key(cur.p, g.cols));
    long long ku = key(cur.p, g.cols);
    if (visited.count(ku)) return;
    visited.insert(ku);
    if (cur.p.r==g.goal.r && cur.p.c==g.goal.c) { fin=true; found=true; path_cache = backtrack(g, parent, cur.p); return; }
    for (auto v : g.neighbors(cur.p)) {
        long long kv = key(v, g.cols);
        if (!visited.count(kv)) {
            parent[kv] = ku;
            pq.push(Node{v, heuristic(g, v, g.goal)});
            frontier_s.insert(kv);
        }
    }
}

AlgoState GBFS::state() const { AlgoState s; s.finished=fin; s.found=found; s.visited=visited; s.frontier=frontier_s; s.path=path_cache; return s; }

void BiBFS::reset(const Grid& g) {
    while(!qa.empty()) qa.pop(); while(!qb.empty()) qb.pop();
    pa.clear(); pb.clear(); va.clear(); vb.clear(); path_cache.clear(); fin=false; found=false; meet.reset();
    qa.push(g.start); va.insert(key(g.start,g.cols));
    qb.push(g.goal); vb.insert(key(g.goal,g.cols));
    frontier_s.clear();
}

void BiBFS::step(const Grid& g) {
    if (fin) return;
    auto expand = [&](std::queue<Pos>& q, std::set<long long>& vset, std::map<long long,long long>& par, const std::set<long long>& other){
        if (q.empty()) return false;
        Pos u = q.front(); q.pop();
        for (auto w : g.neighbors(u)) {
            long long kw = key(w, g.cols);
            if (!vset.count(kw)) {
                vset.insert(kw);
                par[kw] = key(u, g.cols);
                if (other.count(kw)) { meet = w; fin=true; found=true; return true; }
                q.push(w);
                frontier_s.insert(kw);
            }
        }
        return false;
    };
    if (expand(qa, va, pa, vb)) { path_cache = rebuild_bi_path(g, pa, pb, *meet); return; }
    if (expand(qb, vb, pb, va)) { path_cache = rebuild_bi_path(g, pa, pb, *meet); return; }
    if (qa.empty() && qb.empty()) fin=true;
}

AlgoState BiBFS::state() const { AlgoState s; s.finished=fin; s.found=found; s.visited=va; for (auto k:vb) s.visited.insert(k); s.frontier=frontier_s; s.path=path_cache; return s; }
